from __future__ import print_function

import copy
import math
import os
import sys
from collections import defaultdict
from operator import gt
from select import select
from time import sleep

import matplotlib.pyplot as plt
import numpy as np

from KITTI.munkres import Munkres

# import matplotlib; matplotlib.use('Agg')

"""
Data format：
frame_id track_id obj_class  x y z

obj_class:
0 target_car
1 other_car
2 Pestarin
"""

def center_dis(gt,pre):
    delta_x = gt.x - pre.x
    delta_y = gt.y-pre.y
    delta_z = gt.z - pre.z
    return math.sqrt(delta_x*delta_x+ delta_y*delta_y+ delta_z*delta_z)

class tData:
    def __init__(self,frame=-1,obj_class="unset", \
                 x=-1000,y=-1000,z=-1000,track_id=-1):
        # init object data
        self.frame      = frame
        self.track_id   = track_id
        self.obj_class   = obj_class
        self.x          = x
        self.y          = y
        self.z          = z
        self.pre_trj = []

    def __str__(self):
        """
            Print read data.
        """
        attrs = vars(self)
        return '\n'.join("%s: %s" % item for item in attrs.items())

class trackingEvaluation(object):
    """ tracking statistics (CLEAR MOT, id-switches, fragments, ML/PT/MT, precision/recall)
             MOTA   - Multi-object tracking accuracy in [0,100]
             MOTP   - Multi-object tracking precision in [0,100] (3D) / [td,100] (2D)
    """

    def __init__(self, trk_result_path, gt_path, class_type):
        self.threhold = 1
        self.cls = class_type
        gt_file = open(gt_path, "r")
        gt_trks_list = gt_file.readlines()
        end_line_data= np.array(gt_trks_list[len(gt_trks_list)-1].strip().split(' '), dtype = np.float64 )
        self.num_frame = int(end_line_data[0])+1

        self.gt_trajectories            = [[] for x in range(self.num_frame)]
        self.trks_trajectories           = [[] for x in range(self.num_frame)]
        self.loadData(gt_path,True)
        self.loadData(trk_result_path,False)


        self.match_frame = 0
        self.ade = 0
        self.fde = 0

    def loadData(self,path,is_gt=True):
        file = open(path, "r")
        trks_list = file.readlines()
        for i in range(len(trks_list)):
            line_data= np.array(trks_list[i].strip().split(' '), dtype = np.float64 )
            if self.cls == "car" and line_data[2] == 2:
                continue
            if self.cls == "pedestrian" and line_data[2] <=1:
                continue
            t_data  = tData()
            t_data.frame = int(line_data[0])
            t_data.track_id = int(line_data[1])
            t_data.obj_class = int(line_data[2])
            t_data.x = line_data[3]
            t_data.y = line_data[4]
            t_data.z = line_data[5]
            if(is_gt):
                count = 0
                idx = i+1
                pre_temp = []
                while(idx<len(trks_list) and count < 40):
                    line_data = np.array(trks_list[idx].strip().split(' '), dtype = np.float64 )
                    xyz_temp = [0,0,0]
                    if(int(line_data[1]) == t_data.track_id):
                        xyz_temp[0] = line_data[3]
                        xyz_temp[1] = line_data[4]
                        xyz_temp[2] = line_data[5]
                        pre_temp.append(xyz_temp)
                        count+=1
                    idx += 1
                if(len(pre_temp)==40):
                    t_data.pre_trj = np.asarray( pre_temp )
                self.gt_trajectories[t_data.frame].append(t_data)
            else:
                pre_temp = []
                for jj in range(40):
                    xyz_temp = [0,0,0]
                    xyz_temp[0] = line_data[6+jj*3]
                    xyz_temp[1] = line_data[7+jj*3]
                    xyz_temp[2] = line_data[8+jj*3]
                    pre_temp.append(xyz_temp)
                t_data.pre_trj = np.asarray( pre_temp )
                self.trks_trajectories[t_data.frame].append(t_data)
    
    def calMetric(self):
        hm = Munkres()
        max_cost = 1e9

        for frame in  range(self.num_frame):
            is_gt_get_pre = True
            for gg in self.gt_trajectories[frame]:
                if(len(gg.pre_trj)!=40):
                    is_gt_get_pre = False
            if(not is_gt_get_pre):
                continue
            # use hungarian method to associate, using 0..1 as cost
            # build cost matrix
            # row is gt, column is det
            cost_matrix = []
            this_ids_pair = [[],[]]
            for gg in self.gt_trajectories[frame]:
                # save current ids
                this_ids_pair[0].append(gg)
                this_ids_pair[1].append(-1)
                cost_row = []
                for tt in self.trks_trajectories[frame]:
                    c = center_dis(gg, tt)
                    # gating for center distance
                    if c <= self.threhold:
                        cost_row.append(c)
                    else:
                        cost_row.append(max_cost) # = 1e9
                cost_matrix.append(cost_row)

            if len(self.gt_trajectories[frame]) == 0:
                cost_matrix=[[]]
            # associate
            association_matrix = hm.compute(cost_matrix)

            # mapping for tracker ids and ground truth ids
            for row,col in association_matrix:
                # apply gating on box overlap
                c = cost_matrix[row][col]
                if c < max_cost:
                    this_ids_pair[1][row] = self.trks_trajectories[frame][col]

            obj_num = 0
            ade_temp = 0
            fde_temp = 0
            for i in range(len(self.gt_trajectories[frame])):
                if(this_ids_pair[1][i]!=-1):
                    obj_num += 1
                    error = 0
                    for jj in range(40):
                        delta = this_ids_pair[0][i].pre_trj[jj,:] - this_ids_pair[1][i].pre_trj[jj,:]
                        error += np.linalg.norm(delta)
                    error /= 40
                    ade_temp += error
                    delta = this_ids_pair[0][i].pre_trj[39,:] - this_ids_pair[1][i].pre_trj[39,:]
                    f_error = np.linalg.norm(delta)
                    fde_temp += f_error

            if obj_num>0:
                ade_temp /= obj_num
                fde_temp /= obj_num
                self.match_frame += 1
                self.ade += ade_temp
                self.fde += fde_temp

        #MOTA
        ade = self.ade / self.match_frame
        fde = self.fde / self.match_frame
        print("ADE: ",ade)
        print("FDE: ",fde)

if __name__ == "__main__":
    # trk_result_path = "/media/st/data_fast/gazebo/0813/envalute/testing/trajectory/trk_rejuct_tra_result.txt"
    # trk_result_path = "/home/st/ubuntu_data/dataset/gazebo/dynamic_slam/1021_2/testing/trajectory/trk_ls.txt"
    # ADE:  0.5869861244553277
    # FDE:  1.0927926059401352
    # trk_result_path = "/media/st/data_fast/gazebo/0813/envalute/testing/trajectory/trk_without_dis_limit_result.txt"
    # ADE:  0.4700932871477512
    # FDE:  0.8263163793939399
    # trk_result_path = "/media/st/data_fast/gazebo/0813/envalute/testing/trajectory/trk_pc_trk_img_result.txt"
    # gt_path = "/home/st/ubuntu_data/dataset/gazebo/dynamic_slam/1021_2/testing/trajectory/tracking_label.txt"

    # trk_result_path = "/home/st/ubuntu_data/dataset/gazebo/dynamic_slam/0822/testing/trajectory/trj_combine_all.txt"
    # gt_path = "/home/st/ubuntu_data/dataset/gazebo/dynamic_slam/0822/testing/trajectory/tracking_label.txt"

    # trk_result_path = "/home/st/ubuntu_data/dataset/gazebo/dynamic_slam/1021/testing/trajectory/trj_combine_all.txt"
    # gt_path = "/home/st/ubuntu_data/dataset/gazebo/dynamic_slam/1021/testing/trajectory/tracking_label.txt"

    trk_result_path = "/home/st/ubuntu_data/dataset/gazebo/dynamic_slam/1021_2/testing/trajectory/trj_combine_all.txt"
    gt_path = "/home/st/ubuntu_data/dataset/gazebo/dynamic_slam/1021_2/testing/trajectory/tracking_label.txt"
    trk_en = trackingEvaluation(trk_result_path,gt_path,"all")
    trk_en.calMetric()
